from fastapi import APIRouter, Request, Form, status, Query, HTTPException, Body
from fastapi.responses import HTMLResponse, StreamingResponse
from starlette.templating import Jinja2Templates
from models import COLLECTION_NAMES, ConsultationModel, MessageModel
from models.db import get_mongo_db, get_redis_client
from schemas import Consultation, Message
from utils.response import Response
from utils.chatgpt import get_ai_client
from bson import ObjectId
import json
import asyncio
from typing import List, Dict, Any
from datetime import datetime

router = APIRouter(prefix="/api/chat", tags=["问诊"])
templates = Jinja2Templates(directory="templates")


@router.get("/consultations")
async def get_consultations(request: Request):
    """获取用户的问诊列表"""
    user_id = request.state.user["_id"]
    
    # 获取数据库连接实例
    mongo_db = get_mongo_db()

    # 查询用户的所有问诊
    consultations = list(
        mongo_db[COLLECTION_NAMES["consultations"]]
        .find({"user_id": user_id})
        .sort("updated_at", -1)
    )

    # 转换为可序列化的格式
    result = []
    for consultation in consultations:
        # 查询对应的AI角色信息
        role = mongo_db[COLLECTION_NAMES["ai_roles"]].find_one({"_id": consultation["role_id"]})

        result.append({
            "id": str(consultation["_id"]),
            "title": consultation["title"],
            "role_name": role["name"] if role else "未知角色",
            "role_avatar": role["avatar"] if role and "avatar" in role else None,
            "status": consultation["status"],
            "created_at": consultation["created_at"].isoformat(),
            "updated_at": consultation["updated_at"].isoformat()
        })

    return Response.success(data=result)


@router.post("/consultations")
async def create_consultation(request: Request, role_id: str = Body(...), title: str = Body(...)):
    """创建新的问诊会话"""
    user_id = request.state.user["_id"]
    
    # 获取数据库连接实例
    mongo_db = get_mongo_db()
    redis_client = get_redis_client()

    # 验证角色是否存在
    role = mongo_db[COLLECTION_NAMES["ai_roles"]].find_one({"_id": ObjectId(role_id)})
    if not role:
        return Response.error(
            code=status.HTTP_404_NOT_FOUND,
            message="选择的角色不存在"
        )

    # 创建问诊会话
    consultation = ConsultationModel(
        user_id=user_id,
        title=title,
        role_id=ObjectId(role_id)
    )

    # 保存到数据库
    result = mongo_db[COLLECTION_NAMES["consultations"]].insert_one(consultation.to_dict())

    # 初始化Redis中的对话记录
    redis_key = f"consultation:{str(result.inserted_id)}"
    # 添加系统提示词
    system_message = MessageModel(
        role="system",
        content=role["prompt"]
    )
    redis_client.lpush(redis_key, json.dumps(system_message.to_dict()))
    # 设置过期时间（30天）
    redis_client.expire(redis_key, 30 * 24 * 60 * 60)

    return Response.success(
        message="创建成功",
        data={"consultation_id": str(result.inserted_id)}
    )


@router.delete("/consultations/{consultation_id}")
async def delete_consultation(request: Request, consultation_id: str):
    """删除问诊会话"""
    user_id = request.state.user["_id"]
    
    # 获取数据库连接实例
    mongo_db = get_mongo_db()
    redis_client = get_redis_client()

    # 验证问诊是否属于当前用户
    result = mongo_db[COLLECTION_NAMES["consultations"]].delete_one({
        "_id": ObjectId(consultation_id),
        "user_id": user_id
    })

    if result.deleted_count == 0:
        return Response.error(
            code=status.HTTP_404_NOT_FOUND,
            message="问诊会话不存在或无权删除"
        )

    # 删除Redis中的对话记录
    redis_key = f"consultation:{consultation_id}"
    redis_client.delete(redis_key)

    return Response.success(message="删除成功")


@router.put("/consultations/{consultation_id}")
async def update_consultation(request: Request, consultation_id: str, title: str = Body(...)):
    """修改问诊名称"""
    user_id = request.state.user["_id"]
    
    # 获取数据库连接实例
    mongo_db = get_mongo_db()

    # 验证问诊是否属于当前用户
    result = mongo_db[COLLECTION_NAMES["consultations"]].update_one(
        {
            "_id": ObjectId(consultation_id),
            "user_id": user_id
        },
        {
            "$set": {
                "title": title,
                "updated_at": datetime.utcnow()
            }
        }
    )

    if result.matched_count == 0:
        return Response.error(
            code=status.HTTP_404_NOT_FOUND,
            message="问诊会话不存在或无权修改"
        )

    return Response.success(message="修改成功")


@router.get("/consultations/{consultation_id}/messages")
async def get_consultation_messages(request: Request, consultation_id: str):
    """获取问诊会话的消息列表"""
    user_id = request.state.user["_id"]
    
    # 获取数据库连接实例
    mongo_db = get_mongo_db()
    redis_client = get_redis_client()

    # 验证问诊是否属于当前用户
    consultation = mongo_db[COLLECTION_NAMES["consultations"]].find_one({
        "_id": ObjectId(consultation_id),
        "user_id": user_id
    })

    if not consultation:
        return Response.error(
            code=status.HTTP_404_NOT_FOUND,
            message="问诊会话不存在"
        )

    # 从Redis获取消息记录
    redis_key = f"consultation:{consultation_id}"
    messages = redis_client.lrange(redis_key, 0, -1)

    # 转换并过滤掉系统消息
    result = []
    for msg_str in reversed(messages):  # 倒序，使最早的消息在前
        try:
            msg = json.loads(msg_str)
            if msg["role"] != "system":  # 过滤系统消息
                result.append({
                    "role": msg["role"],
                    "content": msg["content"],
                    "timestamp": msg["timestamp"]
                })
        except:
            continue

    return Response.success(data=result)


@router.get("/ai-roles")
async def get_ai_roles(request: Request):
    """获取可用的AI角色列表"""
    # 获取数据库连接实例
    mongo_db = get_mongo_db()
    
    # 查询所有激活的角色
    roles = list(
        mongo_db[COLLECTION_NAMES["ai_roles"]]
        .find({"is_active": True})
        .sort("created_at", 1)
    )

    # 转换为可序列化的格式
    result = []
    for role in roles:
        result.append({
            "id": str(role["_id"]),
            "name": role["name"],
            "description": role.get("description", ""),
            "avatar": role.get("avatar", None)
        })

    return Response.success(data=result)


@router.post("/consultations/{consultation_id}/chat")
async def chat(request: Request, consultation_id: str, message: str = Body(..., embed=True)):
    """发送消息并获取AI回复（流式）"""
    user_id = request.state.user["_id"]
    
    # 获取数据库连接实例
    mongo_db = get_mongo_db()
    redis_client = get_redis_client()

    # 验证问诊是否属于当前用户
    consultation = mongo_db[COLLECTION_NAMES["consultations"]].find_one({
        "_id": ObjectId(consultation_id),
        "user_id": user_id
    })

    if not consultation:
        raise HTTPException(
            status_code=status.HTTP_404_NOT_FOUND,
            detail="问诊会话不存在"
        )

    # 获取角色信息
    role = mongo_db[COLLECTION_NAMES["ai_roles"]].find_one({"_id": consultation["role_id"]})
    if not role:
        raise HTTPException(
            status_code=status.HTTP_404_NOT_FOUND,
            detail="AI角色不存在"
        )

    # 获取供应商信息
    supplier = mongo_db[COLLECTION_NAMES["ai_suppliers"]].find_one({"_id": role["supplier_id"]})
    if not supplier or not supplier.get("is_active", False):
        raise HTTPException(
            status_code=status.HTTP_500_INTERNAL_SERVER_ERROR,
            detail="AI服务不可用"
        )

    # 保存用户消息到Redis
    redis_key = f"consultation:{consultation_id}"
    user_message = MessageModel(role="user", content=message)
    redis_client.lpush(redis_key, json.dumps(user_message.to_dict()))

    # 获取对话历史（包括系统提示词）
    messages_data = redis_client.lrange(redis_key, 0, -1)
    messages = []
    for msg_str in reversed(messages_data):
        try:
            msg = json.loads(msg_str)
            messages.append({
                "role": msg["role"],
                "content": msg["content"]
            })
        except:
            continue
    
    print(f"[对话历史] 发送给AI: {len(messages)}条消息（包含完整历史）")

    # 创建AI客户端
    ai_client = get_ai_client(supplier["name"], supplier["api_key"], supplier["api_url"])
    
    # 从角色配置中获取模型参数
    model = role.get("model")  # 模型名称，如 moonshot-v1-8k
    temperature = role.get("temperature", 0.7)  # 温度参数，默认0.7
    max_tokens = role.get("max_tokens", 1000)  # 最大token数，默认1000
    
    # 调试日志
    print(f"[AI调用] 供应商: {supplier['name']}, 模型: {model}, 温度: {temperature}, 最大Token: {max_tokens}")

    # 异步生成器，用于流式响应
    async def generate():  # 使用异步函数
        full_response = ""
        try:
            # 使用异步生成器获取AI回复，传递模型参数
            async for chunk in ai_client.chat_completion_stream(
                messages, 
                model=model,
                temperature=temperature,
                max_tokens=max_tokens
            ):
                full_response += chunk
                # 返回SSE格式的数据
                yield f"data: {json.dumps({'content': chunk}, ensure_ascii=False)}\n\n"
                await asyncio.sleep(0)

            # 保存AI回复到Redis
            ai_message = MessageModel(role="assistant", content=full_response)
            redis_client.lpush(redis_key, json.dumps(ai_message.to_dict()))

            # 更新问诊会话的更新时间
            mongo_db[COLLECTION_NAMES["consultations"]].update_one(
                {"_id": ObjectId(consultation_id)},
                {"$set": {"updated_at": datetime.utcnow()}}
            )

            # 发送完成信号
            yield f"data: {json.dumps({'done': True})}\n\n"
        except Exception as e:
            # 发送错误信息
            yield f"data: {json.dumps({'error': str(e)})}\n\n"

    # 返回流式响应，设置必要的响应头以禁用缓冲
    return StreamingResponse(
        generate(), 
        media_type="text/event-stream",
        headers={
            "Cache-Control": "no-cache",
            "Connection": "keep-alive",
            "X-Accel-Buffering": "no",  # 禁用 Nginx 缓冲
        }
    )
